import concurrent.futures
import os
import pandas as pd
from xunkemgmt_client.client.api_client import XunkemgmtClient
from cc.appcoin2.strategy.lm.kline_generator2 import convert_response_to_df, generate_request
import json
import time
from collections import defaultdict

os.environ["TZ"] = "GMT"
time.tzset()

def _query_inverval_feed(request_proto):
  try:
    with XunkemgmtClient() as client:
      response = client.query_interval_feed(request_proto)
      return response
  except:
    return None

def get_kline_from_interval_feed(start_dt, end_dt, mea, symbols, return_df=False, resolution="1m", topics=None):
  if topics is None:
    topics = [
        "open_trade",
        "high_trade",
        "low_trade",
        "close_trade",
        "volume",
        "volume_buy",
        "volume_sell",
    ]
    if mea.startswith("Futures"):
      topics.extend(["open_open_interest", "open_funding_rate"])

  if resolution == "1m":
    unit_ns = 60 * 1e9
  elif resolution == "5m":
    unit_ns = 300 * 1e9
  elif resolution == "1h":
    unit_ns = 3600 * 1e9
  elif resolution == "1d":
    unit_ns = 24 * 3600 * 1e9

  print(f"start querying {mea} kline from interval_feed")
  req_list = generate_request(
      mea=mea,
      symbols=symbols,
      start_dt=start_dt,
      end_dt=end_dt,
      resolution=resolution,
      topics=topics
  )

  fs = []
  with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
    for req in req_list:
      future = executor.submit(_query_inverval_feed, req)
      fs.append(future)

  errors = []
  oneday_kline_df = []
  for n in range(len(fs)):
    result = fs[n].result()
    if result is None:
      errors.append("not able to query interval feed")
      continue
    kline_df = convert_response_to_df(result)
    if len(kline_df) != 0:
      oneday_kline_df.append(kline_df)
      if len(result.error_message) > 0:
        errors.append(f"EXCEPTION: {req_list[n].start_time}-{req_list[n].end_time}, {result.error_message}")
    else:
      errors.append(result.error_message)

  if len(errors) > 0:
    for err in errors:
      print(err)

  columns_name_map = {
      "timestamp": "kline_timestamp",
      "open_trade": "open",
      "high_trade": "high",
      "low_trade": "low",
      "close_trade": "close",
      "volume": "volume",
      "volume_buy": "buy_volume",
      "volume_sell": "sell_volume",
      "open_open_interest": "open_interest",
      "open_funding_rate": "funding_rate",
      "market_cap_in_usd": "market_cap_in_usd"
  }

  cols_requiring_value = [
      "kline_timestamp",
      "open",
      "high",
      "low",
      "close",
      "volume"
  ]

  klines = defaultdict(list)
  symbol_kline_df = {}

  if len(oneday_kline_df) == 0:
    return klines

  for symbol in symbols:
    df_list = [df[symbol] for df in oneday_kline_df if symbol in df.keys()]
    if len(df_list) == 0:
      print(f"Interval feed not found for {symbol}, continue")
      continue
    symbol_df = pd.concat(df_list).reset_index(drop=True)
    symbol_df = symbol_df.rename(columns_name_map, axis='columns')
    symbol_df["kline_timestamp"] = (symbol_df["kline_timestamp"] - unit_ns).astype(int)
    if "close" in symbol_df.columns:
      symbol_df["close"] = symbol_df["close"].ffill()
      symbol_df["close"] = symbol_df["close"].bfill()
    for col in ["open", "high", "low"]:
      if col not in symbol_df.columns:
        continue
      mask = symbol_df[col].isnull()
      symbol_df.loc[mask, col] = symbol_df.loc[mask, "close"]
    symbol_df = symbol_df.dropna(subset=[col for col in symbol_df.columns if col in cols_requiring_value]).fillna(0)
    if return_df:
      symbol_kline_df[symbol] = symbol_df
      continue
    symbol_kline = json.loads(symbol_df.to_json(orient="records"))
    klines[symbol] = symbol_kline

  if return_df:
    return symbol_kline_df
  return klines
